import torch
from .vit import VisionTransformer
from data_utils.data_config import DataManager

def build_model_torch(args):
    manager = DataManager(tuple(args.modality_list), args.total_dim, args.seq_len)
    model = VisionTransformer(
        manager=manager,
        in_chans=args.total_dim,
        embed_dim=args.embed_dim,
        depth=args.depth,
        num_heads=args.num_heads,
    )
    return model


def convert_ms_checkpoint(ms_ckpt):
    converted_dict = {}
    for k, v in ms_ckpt.items():
        if not k.startswith('vit'):
            continue
        if 'norm' in k:
            k = k.replace("gamma", "weight")
            k = k.replace("beta", "bias") 
        k = k.replace('vit.', '')
        k = k.replace('cls_tokens', 'cls_token')
        k = k.replace('mask_tokens', 'mask_token')
        converted_dict[k] = torch.from_numpy(v.asnumpy())
    return converted_dict


def convert_ms_checkpoint_notorch(ms_ckpt):
    converted_dict = {}
    for k, v in ms_ckpt.items():
        if not k.startswith('vit'):
            continue
        if 'norm' in k:
            k = k.replace("gamma", "weight")
            k = k.replace("beta", "bias")
        k = k.replace('vit.', '')
        k = k.replace('cls_tokens', 'cls_token')
        k = k.replace('mask_tokens', 'mask_token')
        converted_dict[k] = v.asnumpy()
    return converted_dict


def convert_np_checkpoint_totorch(np_ckpt):
    converted_dict = {}
    for k,v in np_ckpt.items():
        converted_dict[k] = torch.from_numpy(v)
    return converted_dict
